package org.dromara.ai.provider;

import lombok.extern.slf4j.Slf4j;
import org.dromara.ai.domain.dto.AiChatRequest;
import org.dromara.ai.domain.dto.AiChatResponse;
import org.dromara.ai.domain.dto.AiModelInfo;
import org.dromara.ai.domain.enums.AiProviderType;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * AI服务提供商管理器
 * 负责管理所有AI服务提供商的注册、配置和调度
 *
 * @author RuoYi
 * @since 2024-01-01
 */
@Slf4j
@Component
public class AiProviderManager {

    /**
     * 已注册的提供商
     */
    private final Map<String, AiProvider> providers = new ConcurrentHashMap<>();

    /**
     * 提供商类型映射
     */
    private final Map<AiProviderType, List<String>> typeMapping = new ConcurrentHashMap<>();

    /**
     * 模型到提供商的映射
     */
    private final Map<String, String> modelMapping = new ConcurrentHashMap<>();

    /**
     * 默认提供商
     */
    private String defaultProvider;

    /**
     * 注册AI服务提供商
     *
     * @param provider 提供商实例
     */
    public void registerProvider(AiProvider provider) {
        String providerName = provider.getProviderName();
        
        log.info("注册AI服务提供商: {}", providerName);
        
        // 注册提供商
        providers.put(providerName, provider);
        
        // 更新类型映射
        typeMapping.computeIfAbsent(provider.getProviderType(), k -> new ArrayList<>())
                  .add(providerName);
        
        // 更新模型映射
        updateModelMapping(provider);
        
        // 设置默认提供商
        if (defaultProvider == null) {
            defaultProvider = providerName;
        }
        
        log.info("AI服务提供商注册成功: {}", providerName);
    }

    /**
     * 注销AI服务提供商
     *
     * @param providerName 提供商名称
     */
    public void unregisterProvider(String providerName) {
        AiProvider provider = providers.remove(providerName);
        if (provider != null) {
            log.info("注销AI服务提供商: {}", providerName);
            
            // 销毁提供商
            try {
                provider.destroy();
            } catch (Exception e) {
                log.error("销毁AI服务提供商失败: {}", providerName, e);
            }
            
            // 更新类型映射
            typeMapping.get(provider.getProviderType()).remove(providerName);
            
            // 移除模型映射
            removeModelMapping(provider);
            
            // 更新默认提供商
            if (providerName.equals(defaultProvider)) {
                defaultProvider = providers.isEmpty() ? null : providers.keySet().iterator().next();
            }
            
            log.info("AI服务提供商注销成功: {}", providerName);
        }
    }

    /**
     * 获取AI服务提供商
     *
     * @param providerName 提供商名称
     * @return 提供商实例
     */
    public AiProvider getProvider(String providerName) {
        return providers.get(providerName);
    }

    /**
     * 获取默认AI服务提供商
     *
     * @return 默认提供商实例
     */
    public AiProvider getDefaultProvider() {
        return defaultProvider != null ? providers.get(defaultProvider) : null;
    }

    /**
     * 设置默认AI服务提供商
     *
     * @param providerName 提供商名称
     */
    public void setDefaultProvider(String providerName) {
        if (providers.containsKey(providerName)) {
            this.defaultProvider = providerName;
            log.info("设置默认AI服务提供商: {}", providerName);
        } else {
            throw new IllegalArgumentException("提供商不存在: " + providerName);
        }
    }

    /**
     * 根据模型名称获取提供商
     *
     * @param modelName 模型名称
     * @return 提供商实例
     */
    public AiProvider getProviderByModel(String modelName) {
        String providerName = modelMapping.get(modelName);
        return providerName != null ? providers.get(providerName) : null;
    }

    /**
     * 获取所有已注册的提供商
     *
     * @return 提供商列表
     */
    public List<AiProvider> getAllProviders() {
        return new ArrayList<>(providers.values());
    }

    /**
     * 获取可用的提供商
     *
     * @return 可用提供商列表
     */
    public List<AiProvider> getAvailableProviders() {
        return providers.values().stream()
                .filter(AiProvider::isAvailable)
                .collect(Collectors.toList());
    }

    /**
     * 根据类型获取提供商
     *
     * @param type 提供商类型
     * @return 提供商列表
     */
    public List<AiProvider> getProvidersByType(AiProviderType type) {
        return typeMapping.getOrDefault(type, Collections.emptyList()).stream()
                .map(providers::get)
                .filter(Objects::nonNull)
                .collect(Collectors.toList());
    }

    /**
     * 获取所有支持的模型
     *
     * @return 模型列表
     */
    public List<AiModelInfo> getAllSupportedModels() {
        return providers.values().stream()
                .filter(AiProvider::isAvailable)
                .flatMap(provider -> provider.getSupportedModels().stream())
                .collect(Collectors.toList());
    }

    /**
     * 检查模型是否支持
     *
     * @param modelName 模型名称
     * @return 是否支持
     */
    public boolean isModelSupported(String modelName) {
        return modelMapping.containsKey(modelName);
    }

    /**
     * 获取模型信息
     *
     * @param modelName 模型名称
     * @return 模型信息
     */
    public AiModelInfo getModelInfo(String modelName) {
        AiProvider provider = getProviderByModel(modelName);
        return provider != null ? provider.getModelInfo(modelName) : null;
    }

    /**
     * 聊天对话（自动选择提供商）
     *
     * @param request 聊天请求
     * @return 聊天响应
     */
    public AiChatResponse chat(AiChatRequest request) {
        AiProvider provider = selectProvider(request);
        if (provider == null) {
            throw new IllegalStateException("没有可用的AI服务提供商");
        }
        return provider.chat(request);
    }

    /**
     * 异步聊天对话（自动选择提供商）
     *
     * @param request 聊天请求
     * @return 异步聊天响应
     */
    public CompletableFuture<AiChatResponse> chatAsync(AiChatRequest request) {
        AiProvider provider = selectProvider(request);
        if (provider == null) {
            return CompletableFuture.failedFuture(new IllegalStateException("没有可用的AI服务提供商"));
        }
        return provider.chatAsync(request);
    }

    /**
     * 流式聊天对话（自动选择提供商）
     *
     * @param request 聊天请求
     * @param callback 流式回调
     */
    public void chatStream(AiChatRequest request, AiProvider.StreamCallback callback) {
        AiProvider provider = selectProvider(request);
        if (provider == null) {
            callback.onError(new IllegalStateException("没有可用的AI服务提供商"));
            return;
        }
        provider.chatStream(request, callback);
    }

    /**
     * 获取所有提供商统计信息
     *
     * @return 统计信息
     */
    public Map<String, Object> getAllStatistics() {
        Map<String, Object> allStats = new HashMap<>();
        
        allStats.put("totalProviders", providers.size());
        allStats.put("availableProviders", getAvailableProviders().size());
        allStats.put("defaultProvider", defaultProvider);
        allStats.put("totalModels", modelMapping.size());
        
        Map<String, Object> providerStats = new HashMap<>();
        for (Map.Entry<String, AiProvider> entry : providers.entrySet()) {
            providerStats.put(entry.getKey(), entry.getValue().getStatistics());
        }
        allStats.put("providers", providerStats);
        
        return allStats;
    }

    /**
     * 重置所有提供商统计信息
     */
    public void resetAllStatistics() {
        providers.values().forEach(AiProvider::resetStatistics);
        log.info("重置所有AI服务提供商统计信息");
    }

    /**
     * 选择合适的提供商
     *
     * @param request 聊天请求
     * @return 选中的提供商
     */
    private AiProvider selectProvider(AiChatRequest request) {
        // 1. 如果指定了模型，根据模型选择提供商
        if (request.getModel() != null) {
            AiProvider provider = getProviderByModel(request.getModel());
            if (provider != null && provider.isAvailable()) {
                return provider;
            }
        }
        
        // 2. 使用默认提供商
        AiProvider defaultProv = getDefaultProvider();
        if (defaultProv != null && defaultProv.isAvailable()) {
            return defaultProv;
        }
        
        // 3. 选择第一个可用的提供商
        return getAvailableProviders().stream().findFirst().orElse(null);
    }

    /**
     * 更新模型映射
     *
     * @param provider 提供商
     */
    private void updateModelMapping(AiProvider provider) {
        try {
            List<AiModelInfo> models = provider.getSupportedModels();
            for (AiModelInfo model : models) {
                modelMapping.put(model.getId(), provider.getProviderName());
                if (model.getName() != null && !model.getName().equals(model.getId())) {
                    modelMapping.put(model.getName(), provider.getProviderName());
                }
            }
        } catch (Exception e) {
            log.warn("更新模型映射失败: {}", provider.getProviderName(), e);
        }
    }

    /**
     * 移除模型映射
     *
     * @param provider 提供商
     */
    private void removeModelMapping(AiProvider provider) {
        try {
            List<AiModelInfo> models = provider.getSupportedModels();
            for (AiModelInfo model : models) {
                modelMapping.remove(model.getId());
                if (model.getName() != null && !model.getName().equals(model.getId())) {
                    modelMapping.remove(model.getName());
                }
            }
        } catch (Exception e) {
            log.warn("移除模型映射失败: {}", provider.getProviderName(), e);
        }
    }
}